/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */

package SecureUtil;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.SocketTimeoutException;
import java.util.Arrays;
import javax.crypto.Mac;
import javax.crypto.SecretKey;

/**
 *
 * @author Gwang
 */
public class UDPHelper {
    public InetAddress fromip = null;
    public Integer fromport = null;
    DatagramSocket socket;
    public UDPHelper(DatagramSocket socket){
        this.socket = socket;
    }
    public void sendBytes(byte[] tosend, InetAddress ip, int port) throws Exception{
        sendBytes(tosend, ip, port, null);
    }
    public byte[] receiveBytes(int timeout) throws Exception{
        return receiveBytes(timeout, null);
    }
    public void sendBytes(byte[] tosend, InetAddress ip, int port, SecretKey sk) throws Exception{
        byte[] hmac = null;
        if(sk!=null){
            Mac macer = Mac.getInstance("HmacMD5");
            macer.init(sk);
            hmac = macer.doFinal(tosend);
        }
        Packer packer = new Packer();
        packer.addElement(Packer.DataType.BYTES, tosend);
        if(hmac!=null)
            packer.addElement(Packer.DataType.BYTES, hmac);
        byte[] bytes = packer.pack();
        DatagramPacket packet = new DatagramPacket(bytes,bytes.length,ip,port);
        socket.send(packet);
    }
    public byte[] receiveBytes(int timeout, SecretKey sk)
    throws Exception{
        socket.setSoTimeout(timeout);
               
        byte[] buf = new byte[65507];
        DatagramPacket packet = new DatagramPacket(buf,buf.length);
        byte[] msg = null;
        byte[] hmacback = null;
        try{
            socket.receive(packet);
            fromip = packet.getAddress();
            fromport = packet.getPort();
            Packer packer = new Packer();
            packer.unpack(packet.getData());
            try{
                msg = (byte[])packer.getNext(Packer.DataType.BYTES);
                if(sk != null)
                    hmacback = (byte[])packer.getNext(Packer.DataType.BYTES);
            }catch(Exception e){
                e.printStackTrace();
                return null;
            }
            if(sk!=null){
                Mac macer = Mac.getInstance("HmacMD5");
                macer.init(sk);
                byte[] hmac = macer.doFinal(msg);
                if(!Arrays.equals(hmac, hmacback)){
                   throw new Exception("integrity dangerous: hmac does not match");
                }
            }
        }catch(SocketTimeoutException e){
            System.out.println("socket time out; probabaly being DoS attacked");
            return null;
        }catch(IOException e){
            e.printStackTrace();
        }
        return msg;
    }
}
